#include <iostream>
#include <opencv2/opencv.hpp>

class CameraParam
{
public:
    CameraParam(const std::string &config_path);
    ~CameraParam(){};
    cv::Mat imageUndist(const cv::Mat& src);

    double fx_, fy_, cx_, cy_;
    double k1_, k2_, p1_, p2_, k3_;
    int w_, h_;
    std::string config_path_;

    cv::Mat cam_matrix_, cam_dist_;
};

CameraParam::CameraParam(const std::string &config_path) : config_path_(config_path)
{
    cv::FileStorage fs(config_path_, cv::FileStorage::READ);
    if(!fs.isOpened())
    {
        std::cout << "\033[1;31m" << "Open " +  config_path_ + " is failure..." << "\033[0m\n";
        return;
    }
    cv::FileNode n =  fs["projection_parameters"];
    fx_ = static_cast<double> (n["fx"]);
    fy_ = static_cast<double> (n["fy"]);
    cx_ = static_cast<double> (n["cx"]);
    cy_ = static_cast<double> (n["cy"]);

    n = fs["distortion_parameters"];
    k1_ = static_cast<double> (n["k1"]);
    k2_ = static_cast<double> (n["k2"]);
    p1_ = static_cast<double> (n["p1"]);
    p2_ = static_cast<double> (n["p2"]);
    k3_ = static_cast<double> (n["k3"]);

    w_ = fs["image_width"];
    fs["image_height"] >> h_;

    cam_matrix_ = (cv::Mat_<float>(3, 3) << fx_, 0,  cx_, 
                                           0,  fy_, cy_,
                                           0,  0,   1);
    
    cam_dist_ = (cv::Mat_<float>(5, 1) << k1_, k2_, p1_, p2_, k3_);
    
    std::cout <<  "\033[1;32m \nimage size: " << w_ << " x " << h_ << std::endl;
    std::cout << "projection parameters:\n" << cam_matrix_ << std::endl;
    std::cout << "distortion parameters:\n" << cam_dist_.t() << std::endl << std::endl;
}

cv::Mat CameraParam::imageUndist(const cv::Mat &img)
{
    cv::Mat new_camera_matrix;
    cv::Rect roi;
    getOptimalNewCameraMatrix(cam_matrix_,
                              cam_dist_,
                              cv::Size(w_,h_),
                              1,
                              cv::Size(w_,h_),
                            //   &new_camera_matrix,
                              &roi);

    // Undistort the image
    cv::Mat dst;
    undistort(img,
              dst,
              cam_matrix_,
              cam_dist_);

    // Crop the image
    dst = dst(roi);
    return dst;
}
